import numpy as np
import matplotlib.pyplot as plt
from investment.time_series import IndexTimeSeries, filter_indexes_connection
from investment.date_utilities import DateUtility
from investment.operation import calc_reward_rate


def put_into_bins(series, bin_count):
    """
    :param series: array of series
    :param bin_count: int like 30
    :return: Tuple[float array of lin_space by min-max-num_of_bins, bin box List of each index which split by step]
    """
    series_bins = np.linspace(series.min(), series.max(), bin_count)
    series_put_bins = np.digitize(series, series_bins)
    bin_index_box = []
    for b in range(1, bin_count + 1):
        bin_where = np.where(series_put_bins == b)[0]
        bin_index_box.append(bin_where)

    return series_bins, bin_index_box


def plot_heat_map(matrix, series_y_ticks):
    heap_mat = matrix.copy()
    heap_mat[heap_mat == 0] = np.nan
    fig = plt.figure()
    ax = plt.gca()
    im = ax.imshow(np.flip(heap_mat, 0), aspect='auto', cmap='jet')  # aspect auto to square!
    plt.yticks(range(heap_mat.shape[0]), np.around(series_y_ticks[::-1], decimals=2))
    plt.xticks(range(0, heap_mat.shape[1], 12), [f"{i}yr" for i in range(10)])
    fig.colorbar(im, fraction=0.046, pad=0.04)
    plt.show()


class DiamondMatrix:
    def __init__(self, evaluation_series, price_series, date_utility, evaluation_bins=30, max_year=10):
        assert isinstance(date_utility, DateUtility)
        self.evaluation_series = evaluation_series
        self.price_series = price_series
        self.date_utility = date_utility
        self.evaluation_bins = evaluation_bins
        self.month_counts = max_year * 12
        self.evaluation_space: np.ndarray = None
        self.evaluation_index_buckets: list = None
        self.evaluation_space, self.evaluation_index_buckets = put_into_bins(evaluation_series, evaluation_bins)
        # reward_matrix to save all reward results
        self.reward_matrix: np.ndarray = None
        self.shape = [self.evaluation_bins, self.month_counts]

    def _calc_reward_matrix(self):
        reward_matrix = [[[] for _ in range(self.month_counts)] for _ in range(self.evaluation_bins)]
        for i_pe, pe in enumerate(self.evaluation_space):
            for month in range(1, self.month_counts + 1):
                for pe_index in self.evaluation_index_buckets[i_pe]:
                    find_idx = self.date_utility.get_interval_trading_index(pe_index, month, d_format="m")
                    if find_idx is not None:
                        reward = calc_reward_rate(self.price_series[pe_index], self.price_series[find_idx])
                        reward_matrix[i_pe][month - 1].append(reward)
        self.reward_matrix = reward_matrix
        return self.reward_matrix

    def get_statistic_matrix(self, func=np.max):
        if self.reward_matrix is None:
            self._calc_reward_matrix()
        matrix = np.zeros([self.evaluation_bins, self.month_counts])
        # a = np.array(map(lambda x: map(lambda y: np.max(y, initial=np.nan)), reward_matrix))
        for i_pe, pe in enumerate(self.evaluation_space):
            for month in range(1, self.month_counts + 1):
                reward_records = self.reward_matrix[i_pe][month - 1]
                matrix[i_pe, month - 1] = func(reward_records) if len(reward_records) > 0 else np.nan
        return matrix

    def get_happen_time_matrix(self):
        time_matrix = np.zeros([self.evaluation_bins, self.month_counts])
        for i_pe, pe in enumerate(self.evaluation_space):
            for month in range(1, self.month_counts + 1):
                count = 0
                for pe_index in self.evaluation_index_buckets[i_pe]:
                    find_date = self.date_utility.get_interval_trading_index(pe_index, month, d_format="m")
                    if find_date is not None:
                        count += 1
                time_matrix[i_pe, month - 1] = count
        return time_matrix

    def which_bins(self, value) -> int:
        """
        which evaluation bins the given value belongs to
        :param value: evaluation value
        :return: int: index of evaluation bins array
        """
        return np.nansum(self.evaluation_space < value) - 1

    def get_evaluation_bucket(self, value) -> np.ndarray:
        indexes = self.evaluation_index_buckets[self.which_bins(value)]
        return filter_indexes_connection(indexes, pad=10)


def tt():
    import timeit
    mysetup = "from __main__ import search_trade_day, trading_dates, add_day"
    mycode = "search_trade_day(trading_dates, add_day)"
    print(timeit.timeit(setup=mysetup, stmt=mycode, number=100))


if __name__ == "__main__":
    from investment.constants import A_INDEX_DICT
    from quantify.value_time_series import plot_index_fundamental
    import os
    plt.ion()
    a_market = IndexTimeSeries(A_INDEX_DICT['沪深300'], "2002-07-01", "2022-07-01", fill_flag=True)
    # a_market = IndexTimeSeries(A_INDEX_DICT['沪深300'], "1997-11-01", "2013-04-01", fill_flag=True)
    pe_series = a_market.get_fundamental_metrics("PE")
    cp_series = a_market.get_fundamental_metrics("CP")
    plot_index_fundamental(a_market, "PE")
    diamond_mat = DiamondMatrix(pe_series, cp_series, DateUtility(a_market.dates_series))
    # happen = diamond_mat.get_happen_time_matrix()
    # plot_heat_map(happen, diamond_mat.series_bins)
    for fun in [np.max, np.min, np.median]:
        mat = diamond_mat.get_statistic_matrix(func=fun)
        plot_heat_map(mat, diamond_mat.evaluation_space)
        with open(os.path.join("../test/data", f"mat_{str(fun.__name__)}.npy"), 'wb') as f:
            np.save(f, mat)
    print("f")
